import cv2
import numpy as np
import yaml
from collections import defaultdict
from discret_image import active_size_map_discretization


MODES = ["turn", "node_edge_mode", "auto_mode", "visulize"]
mode = MODES[0]

pgm_path = "../maps/map.pgm"
config_path = "../maps/map.yaml"
png_path = pgm_path.replace(".pgm", ".png")
graph_path = pgm_path.replace(".pgm", "_graph.yaml")

# 离散化的细粒度半径，越小越细，节点越多
NODE_RADIUS = 50


def turn_pgm_to_png(pgm_path, png_path):
    img = cv2.imread(pgm_path, -1)
    if img is None:
        raise ValueError(f"无法读取PGM文件: {pgm_path}")

    if len(img.shape) == 2:
        img_rgb = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)
    else:
        img_rgb = img[:, :, :3].copy()  # 避免内存重叠问题

    cv2.imwrite(png_path, img_rgb)


def draw_graph_on_image(img, graph):
    nodes = np.array(graph["nodes"], dtype=np.int32)
    edges = graph["edges"]

    img_with_graph = cv2.cvtColor(
        cv2.cvtColor(img, cv2.COLOR_BGR2GRAY), cv2.COLOR_GRAY2BGR
    )

    for idx in range(len(nodes)):
        cv2.circle(img_with_graph, tuple(nodes[idx][:2]), 1, (0, 255, 0), -1)

    for e in edges:
        cv2.line(
            img_with_graph,
            tuple(nodes[e[0]][:2]),
            tuple(nodes[e[1]][:2]),
            (255, 0, 255),
            1,
        )

    cv2.imshow("Graph", img_with_graph)
    cv2.waitKey(0)
    cv2.destroyAllWindows()
    return img_with_graph


def visualize_graph(img_path, graph_path, config_path):
    img = cv2.imread(img_path)
    with open(graph_path, "r") as f:
        graph = yaml.safe_load(f)
    with open(config_path, "r") as f:
        config = yaml.safe_load(f)
    resolution = config["resolution"]
    x, y, yaw = config["origin"]
    yaw = np.deg2rad(yaw)
    R = np.array([[np.cos(yaw), -np.sin(yaw)], [np.sin(yaw), np.cos(yaw)]])
    T = np.array([x, y])

    # 世界坐标 → 图像坐标（带Y轴翻转）
    height, width = img.shape[:2]
    transformed_nodes = []
    for node in graph["nodes"]:
        world = np.array(node[:2])
        img_coords = R.T @ ((world - T) / resolution)
        x_img, y_img = img_coords
        y_img = height - y_img  # Y轴翻转
        transformed_nodes.append([x_img, y_img])
    graph["nodes"] = transformed_nodes
    draw_graph_on_image(img, graph)


def save_graph_to_yaml(graph, config_path, graph_path):
    with open(config_path, "r") as f:
        config = yaml.safe_load(f)
    resolution = config["resolution"]
    x, y, yaw = config["origin"]
    yaw = np.deg2rad(yaw)
    R = np.array([[np.cos(yaw), -np.sin(yaw)], [np.sin(yaw), np.cos(yaw)]])
    T = np.array([x, y])

    # 图像坐标 → 世界坐标（带Y轴翻转）
    height, width = cv2.imread(png_path).shape[:2]
    transformed_nodes = []
    for node in graph["nodes"]:
        x_img, y_img = node[:2]
        y_img = height - y_img  # Y轴翻转
        world_coords = R @ (np.array([x_img, y_img]) * resolution) + T
        transformed_nodes.append(world_coords.tolist())
    graph["nodes"] = [[x, y, 0] for x, y in transformed_nodes]

    with open(graph_path, "w") as f:
        yaml.dump(graph, f, default_flow_style=False)


def get_graph_from_image(png_path, config_path, graph_path):
    """
    Extracts a graph from an image where:
    - Nodes are red dots [0, 0, 255] in BGR
    - Edges are blue lines [255, 0, 0] in BGR
    - Converts the graph to a real-world coordinate system using config.
    """
    # Load image
    img = cv2.imread(png_path)

    # --- Step 1: Node Detection (Red Dots) ---
    lower_red = np.array([0, 0, 200])
    upper_red = np.array([0, 0, 255])
    mask_red = cv2.inRange(img, lower_red, upper_red)

    # Find contours of red regions
    contours_red, _ = cv2.findContours(
        mask_red, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
    )
    nodes = []
    for cnt in contours_red:
        M = cv2.moments(cnt)
        if M["m00"] != 0:
            x = int(M["m10"] / M["m00"])
            y = int(M["m01"] / M["m00"])
            nodes.append([x, y])

    # --- Step 2: Edge Detection (Blue Lines) ---
    lower_blue = np.array([200, 0, 0])
    upper_blue = np.array([255, 0, 0])
    mask_blue = cv2.inRange(img, lower_blue, upper_blue)

    # Morphological closing to connect broken lines
    kernel = np.ones((3, 3), np.uint8)
    mask_blue = cv2.morphologyEx(mask_blue, cv2.MORPH_CLOSE, kernel)

    # Find all blue pixels
    points = cv2.findNonZero(mask_blue)
    if points is None or len(points) < 10:
        edges = []
    else:
        points = points.reshape(-1, 2).astype(np.float32)

        # 分割连通区域
        num_labels, labels = cv2.connectedComponents(mask_blue)

        # 将点按照连通区域标签分组
        label_to_points = defaultdict(list)
        for point in points:
            x, y = int(point[0]), int(point[1])
            label = labels[y, x]
            label_to_points[label].append([x, y])

        edges = []

        # 对每个连通区域单独处理
        for label in range(1, num_labels):  # 跳过背景（label=0）
            cluster_points = np.array(label_to_points[label], dtype=np.float32)

            if len(cluster_points) < 2:
                continue

            # 使用PCA拟合直线方向
            mean, eig, _ = cv2.PCACompute2(cluster_points, mean=np.array([]))
            direction = eig[0]

            # 投影到主方向上
            vectors = cluster_points - mean[0]
            projections = np.dot(vectors, direction)

            min_idx = np.argmin(projections)
            max_idx = np.argmax(projections)
            pt1 = tuple(cluster_points[min_idx].astype(int))
            pt2 = tuple(cluster_points[max_idx].astype(int))

            # 寻找最近的两个节点
            min_dist1 = float("inf")
            min_idx1 = -1
            min_dist2 = float("inf")
            min_idx2 = -1

            for idx, node in enumerate(nodes):
                dist1 = np.linalg.norm(np.array(pt1) - np.array(node))
                dist2 = np.linalg.norm(np.array(pt2) - np.array(node))

                if dist1 < min_dist1:
                    min_dist1 = dist1
                    min_idx1 = idx
                if dist2 < min_dist2:
                    min_dist2 = dist2
                    min_idx2 = idx

            edges.append([min_idx1, min_idx2])

    # --- Step 3: Build Graph Structure ---
    graph = {
        "nodes": nodes,
        "edges": edges,
    }
    draw_graph_on_image(img, graph)

    # 保存
    save_graph_to_yaml(graph, config_path, graph_path)
    return graph


def get_graph_by_discretization(img_path, config_path, graph_path, root=[0, 0]):
    """
    通过图像离散化生成图结构，包含图像预处理步骤
    """
    # 1. 读取图像
    img = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE)

    # 2. 图像预处理
    print("[预处理] 开始图像降噪和增强...")
    img_denoised = cv2.medianBlur(img, 3)
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
    img_closed = cv2.morphologyEx(img_denoised, cv2.MORPH_CLOSE, kernel, iterations=1)
    img_smooth = cv2.GaussianBlur(img_closed, (3, 3), 0)
    _, binary_img = cv2.threshold(img_smooth, 150, 255, cv2.THRESH_BINARY)

    # 3. 读取配置参数
    with open(config_path, "r") as f:
        config = yaml.safe_load(f)
    resolution = config["resolution"]
    x, y, yaw = config["origin"]
    # ROS map_saver: yaw为弧度，像素坐标原点在左下角，y轴向上
    # 图像坐标系：左上角为(0,0)，x向右，y向下
    height, width = img.shape[:2]
    # 世界坐标转像素坐标
    root_world = np.array(root[:2], dtype=np.float32)
    # 先减去origin平移，再旋转
    R = np.array(
        [[np.cos(yaw), -np.sin(yaw)], [np.sin(yaw), np.cos(yaw)]], dtype=np.float32
    )
    T = np.array([x, y], dtype=np.float32)
    # 世界坐标转到地图坐标
    map_xy = (R.T @ (root_world - T)) / resolution
    # 地图坐标转到像素坐标（y轴翻转）
    root_pixel = np.array([map_xy[0], height - 1 - map_xy[1]])
    root_pixel = np.round(root_pixel).astype(int)
    root_pixel[0] = np.clip(root_pixel[0], 0, img.shape[1] - 1)
    root_pixel[1] = np.clip(root_pixel[1], 0, img.shape[0] - 1)
    root_pixel = tuple(root_pixel)
    
    # 5. 使用漫水填充标记可行区域
    h, w = img.shape
    mask = np.zeros((h + 2, w + 2), dtype=np.uint8)
    flooded = binary_img.copy()
    flood_color = 100
    lo_diff, hi_diff = 10, 10
    flags = 4 | (lo_diff << 8) | cv2.FLOODFILL_FIXED_RANGE
    cv2.floodFill(flooded, mask, root_pixel, flood_color, lo_diff, hi_diff, flags)

    # 6. 调用离散化函数
    passable_color = flood_color
    result_img, graph = active_size_map_discretization(
        flooded, passable_color=passable_color, node_radius=NODE_RADIUS
    )

    # 7. 【新增】统计节点度数并删除孤立点
    print("[新增] 正在统计节点度数并删除孤立点...")
    nodes = graph["nodes"]
    edges = graph["edges"]

    # 7.1 统计度数
    degrees = [0] * len(nodes)
    for u, v in edges:
        degrees[u] += 1
        degrees[v] += 1

    # 7.2 找出孤立点
    isolated_nodes = set(i for i, d in enumerate(degrees) if d == 0)
    print(f"  >> 孤立点数量: {len(isolated_nodes)}")

    # 7.3 构建新节点列表和索引映射
    new_nodes = []
    old_to_new = {}
    for old_idx, node in enumerate(nodes):
        if old_idx not in isolated_nodes:
            old_to_new[old_idx] = len(new_nodes)
            new_nodes.append(node)

    # 7.4 构建新边列表
    new_edges = []
    for u, v in edges:
        if u not in isolated_nodes and v not in isolated_nodes:
            new_u = old_to_new[u]
            new_v = old_to_new[v]
            new_edges.append([new_u, new_v])

    # 7.5 更新图结构
    graph["nodes"] = new_nodes
    graph["edges"] = new_edges
    print(f"  >> 更新后节点数: {len(new_nodes)}, 边数: {len(new_edges)}")

    # 8. 可视化结果
    cv2.imwrite("../maps/temple.png", result_img)

    # 9. 保存图结构
    save_graph_to_yaml(graph, config_path, graph_path)
    return graph


if __name__ == "__main__":
    print("Running mode:", mode)
    if mode == MODES[0]:
        turn_pgm_to_png(pgm_path, png_path)
    elif mode == MODES[1]:
        get_graph_from_image(png_path, config_path, graph_path)
    elif mode == MODES[2]:
        get_graph_by_discretization(png_path, config_path, graph_path)
    else:
        visualize_graph(png_path, graph_path, config_path)
